#include <boost/asio.hpp>
#include <nlohmann/json.hpp>
#include <iostream>
#include <memory>
#include <string>

using boost::asio::ip::tcp;
using json = nlohmann::json;

class Session : public std::enable_shared_from_this<Session> {
public:
    Session(tcp::socket socket) : socket_(std::move(socket)) {}

    void start() {
        auto self(shared_from_this());
        // Send hello message to client
        json hello_msg = {
            {"type", "hello"},
            {"message", "Hello from server!"}
        };
        std::string msg = hello_msg.dump() + "\n";  // Add newline as message delimiter
        
        boost::asio::async_write(socket_, boost::asio::buffer(msg),
            [this, self](boost::system::error_code ec, std::size_t /*length*/) {
                if (!ec) {
                    do_read();
                }
            });
    }

private:
    void do_read() {
        auto self(shared_from_this());
        boost::asio::async_read_until(socket_, data_, '\n',
            [this, self](boost::system::error_code ec, std::size_t length) {
                if (!ec) {
                    std::string message(boost::asio::buffers_begin(data_.data()),
                                      boost::asio::buffers_begin(data_.data()) + length);
                    data_.consume(length);

                    try {
                        auto j = json::parse(message);
                        std::cout << "Received from client: " << j.dump(2) << std::endl;
                    } catch (const json::parse_error& e) {
                        std::cerr << "JSON parse error: " << e.what() << std::endl;
                    }

                    do_read();  // Continue reading
                }
            });
    }

    tcp::socket socket_;
    boost::asio::streambuf data_;
};

class Server {
public:
    Server(boost::asio::io_context& io_context, short port)
        : acceptor_(io_context, tcp::endpoint(tcp::v4(), port)) {
        do_accept();
    }

private:
    void do_accept() {
        acceptor_.async_accept(
            [this](boost::system::error_code ec, tcp::socket socket) {
                if (!ec) {
                    std::cout << "New client connected!" << std::endl;
                    std::make_shared<Session>(std::move(socket))->start();
                }

                do_accept();
            });
    }

    tcp::acceptor acceptor_;
};

int main() {
    try {
        boost::asio::io_context io_context;
        Server server(io_context, 12345);
        std::cout << "Server started on port 12345" << std::endl;
        io_context.run();
    } catch (std::exception& e) {
        std::cerr << "Exception: " << e.what() << std::endl;
    }

    return 0;
}